import os
import shutil
import argparse
import cv2
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image
import satellite_cloud_generator as scg
import torch


parser = argparse.ArgumentParser()

parser.add_argument('--interval', type=int, default=3)
parser.add_argument('--light', type=int, default=0)
parser.add_argument('--fog', type=int, default=0)
parser.add_argument('--gaussian', type=int, default=0)
parser.add_argument('--white', type=int, default=0)
parser.add_argument('--blur', type=int, default=0)
parser.add_argument('--geo', type=int, default=0)
args = parser.parse_args()

input_folder = os.path.join(os.getcwd(), 'input_images')
output_folder = os.path.join(os.getcwd(), 'output_images')
if not os.path.exists(input_folder):
    os.makedirs(input_folder)
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

del_list = os.listdir(output_folder)
for f in del_list:
    file_path = os.path.join(output_folder, f)
    if os.path.isfile(file_path):
        os.remove(file_path)

input_images = os.listdir(input_folder)

for index, image_file in enumerate(input_images):
    if index % args.interval != 0:
        continue
    input_path = os.path.join(input_folder, image_file)
    output_path = os.path.join(output_folder, image_file)
    # 调整亮度
    image = cv2.imread(input_path)
    if args.light >= 1 and args.light <= 99:
        # 调整亮度，对比度系数设置为1表示不改变对比度
        image = cv2.convertScaleAbs(image, alpha=1, beta=args.light)
        # cv2.imwrite(output_path, image)
    # 添加白噪声
    if args.white >= 1 and args.white <= 99:
        # 读取原图像
        # image = cv2.imread(input_path)
        # 获取图像的高度和宽度
        height, width, channels = image.shape
        # 生成与原图像相同尺寸的白噪声
        noise = np.random.normal(0, args.white / 50,
                                 (height, width, channels)).astype(np.uint8)
        # 将噪声添加到原图像上
        # noisy_image = cv2.add(image, noise)
        image = cv2.add(image, noise)
        # 保存加噪后的图像
        # cv2.imwrite(output_path, noisy_image)
        # 添加高斯模糊
    if args.blur >= 1 and args.blur <= 99:
        # image = cv2.imread(input_path)
        # 如果blur_strength是偶数，将其调整为最接近的奇数
        if args.blur % 2 == 0:
            args.blur += 1
        # 添加高斯模糊
        # blurred_image = cv2.GaussianBlur(
        #     image, (args.blur, args.blur), 0)
        image = cv2.GaussianBlur(
            image, (args.blur, args.blur), 0)
        # 保存加噪后的图像
    if args.geo >= 1 and args.geo <= 100:
        w, h = image.shape[1], image.shape[0]
        k = np.zeros((3,3))
        k[0, 0] = w
        k[1, 1] = h
        k[0, 2] = w / 2.0
        k[1, 2] = h / 2.0
        k[2, 2] = 1.0
        dist_coef = np.zeros((4,1))
        dist_coef[0,0] = args.geo / 100
        image = cv2.undistort(image, k, dist_coef)
    cv2.imwrite(output_path, image)
    # 添加高斯噪声
    if args.gaussian >= 1 and args.gaussian <= 99:
        # image = Image.open(input_path)
        image = Image.open(output_path)
        # 将图像转换为NumPy数组
        image_np = np.array(image)
        # 生成与图像相同形状的高斯噪声
        mean = 0
        std = args.gaussian
        noise = np.random.normal(mean, std, image_np.shape).astype(np.uint8)
        # 添加噪声到图像
        noisy_image_np = np.clip((image_np + noise), 0, 255).astype(np.uint8)
        # 创建添加噪声后的图像对象
        image = Image.fromarray(noisy_image_np)
        # 保存添加噪声后的图像
        image.save(output_path)
    # 添加雾
    if args.fog >= 1 and args.fog <= 99:
        image = np.array(Image.open(output_path))
        image_tensor = torch.from_numpy(np.transpose(image, (2, 0, 1))).unsqueeze(0).float() / 255.0
        cloudy_image, cloud_mask = scg.add_cloud(image_tensor, min_lvl=0.0, max_lvl=args.fog / 99, return_cloud=True)
        result_image = np.transpose(cloudy_image.squeeze().numpy(), (1, 2, 0))
        im = Image.fromarray(np.uint8(result_image * 255))
        im.save(output_path)